# notes.py
import json
import os
import datetime
from datetime import datetime
import markdown
from PyQt5.QtWidgets import *
from PyQt5.QtCore import *
from PyQt5.QtGui import *
from PyQt5.QtWebEngineWidgets import QWebEngineView

NOTES_FILE = "data/notes_data.json"
NOTES_DIR = "data/notes"

class NotesModule(QWidget):
    content_updated = pyqtSignal(dict)
    
    def __init__(self):
        super().__init__()
        self.notes = []
        self.folders = []
        self.tags = ["所有标签", "无标签"]
        self.current_note = None
        self.expanded_paths = set()
        
        os.makedirs(NOTES_DIR, exist_ok=True)
        self.init_ui()
        self.load_data()

    def init_ui(self):
        main_layout = QHBoxLayout(self)
        main_layout.setContentsMargins(15, 15, 15, 15)
        main_layout.setSpacing(20)

        # 左侧导航面板
        left_panel = self.create_left_panel()
        
        # 右侧编辑预览区
        right_panel = self.create_right_panel()
        
        main_layout.addWidget(left_panel)
        main_layout.addWidget(right_panel)
        self.setStyleSheet("""
            QTreeWidget, QListWidget {
                background: white;
                border-radius: 8px;
                border: 1px solid #e0e0e0;
            }
            QLineEdit, QTextEdit {
                padding: 8px;
                border: 1px solid #e0e0e0;
                border-radius: 8px;
            }
        """)

    def create_left_panel(self):
        panel = QWidget()
        panel.setFixedWidth(300)
        layout = QVBoxLayout(panel)
        
        self.mode_tabs = QTabWidget()
        self.folder_tree = QTreeWidget()
        self.folder_tree.setHeaderHidden(True)
        self.folder_tree.itemClicked.connect(self.load_note)
        self.folder_tree.setContextMenuPolicy(Qt.CustomContextMenu)
        self.folder_tree.customContextMenuRequested.connect(self.show_tree_context_menu)
        
        self.tag_list = QListWidget()
        self.tag_list.addItem("无标签")
        self.tag_list.itemClicked.connect(self.filter_by_tag)
        self.tag_list.setContextMenuPolicy(Qt.CustomContextMenu)
        self.tag_list.customContextMenuRequested.connect(self.show_tag_context_menu)
        
        self.mode_tabs.addTab(self.folder_tree, "文件夹")
        self.mode_tabs.addTab(self.tag_list, "标签筛选")
        
        self.new_btn = QPushButton("新建笔记")
        self.new_btn.clicked.connect(self.create_note)
        self.new_btn.setStyleSheet("""
            QPushButton {
                background: #2196F3;
                color: white;
                padding: 12px;
                border-radius: 8px;
                font-weight: 500;
            }
            QPushButton:hover { background: #1976D2; }
        """)
        
        layout.addWidget(self.mode_tabs)
        layout.addWidget(self.new_btn)
        return panel

    def create_right_panel(self):
        panel = QWidget()
        layout = QVBoxLayout(panel)
        
        # 标题和标签水平布局
        header_layout = QHBoxLayout()
        self.title_input = QLineEdit()
        self.title_input.setPlaceholderText("输入标题...")
        self.tag_input = QLineEdit()
        self.tag_input.setPlaceholderText("输入标签（逗号分隔）...")
        header_layout.addWidget(self.title_input, 3)
        header_layout.addWidget(self.tag_input, 2)
        
        # 统一标题栏
        titles_layout = QHBoxLayout()
        editor_title = QLabel("编辑区（支持Markdown + LaTeX）")
        preview_title = QLabel("预览区")
        for title in [editor_title, preview_title]:
            title.setStyleSheet("""
                QLabel {
                    font: bold 14px 'Segoe UI';
                    color: #2c3e50;
                    padding: 8px 0;
                }
            """)
        titles_layout.addWidget(editor_title, 1)
        titles_layout.addWidget(preview_title, 1)
        
        # 左右分栏布局
        splitter = QSplitter(Qt.Horizontal)
        
        # 左侧编辑区
        editor_box = QWidget()
        editor_layout = QVBoxLayout(editor_box)
        editor_layout.setContentsMargins(0, 0, 0, 0)
        self.editor = QTextEdit()
        self.editor.textChanged.connect(self.update_preview)
        editor_layout.addWidget(self.editor)
        
        # 右侧预览区
        preview_box = QWidget()
        preview_layout = QVBoxLayout(preview_box)
        preview_layout.setContentsMargins(0, 0, 0, 0)
        self.preview = QWebEngineView()
        preview_layout.addWidget(self.preview)
        
        splitter.addWidget(editor_box)
        splitter.addWidget(preview_box)
        splitter.setSizes([500, 500])
        
        # 保存按钮
        self.save_btn = QPushButton("保存")
        self.save_btn.clicked.connect(self.save_current)
        self.save_btn.setStyleSheet("""
            QPushButton {
                background: #4CAF50;
                color: white;
                padding: 12px 24px;
                border-radius: 8px;
            }
            QPushButton:hover { background: #45a049; }
        """)
        
        layout.addLayout(header_layout)
        layout.addLayout(titles_layout)
        layout.addWidget(splitter, 1)
        layout.addWidget(self.save_btn)
        return panel

    def update_preview(self):
        content = self.editor.toPlainText()
        html = markdown.markdown(content, extensions=['extra', 'codehilite'])
        
        mathjax_config = """
        <script>
        MathJax = {
          tex: {
            inlineMath: [['$', '$'], ['\\(', '\\)']],
            displayMath: [['$$', '$$'], ['\\[', '\\]']],
            processEscapes: true,
            packages: {'[+]': ['amsmath']}
          },
          options: {
            ignoreHtmlClass: 'tex2jax_ignore',
            processHtmlClass: 'tex2jax_process'
          },
          loader: {load: ['[tex]/amsmath']},
          startup: {
            ready: () => {
              MathJax.startup.defaultReady();
              MathJax.startup.promise.then(() => {
                console.log('MathJax initial typesetting complete');
              });
            }
          }
        };
        </script>
        <script src="https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js" async></script>
        """
        
        self.preview.setHtml(f"""
            <!DOCTYPE html>
            <html>
            <head>
                <meta charset="utf-8">
                <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.5.0/styles/default.min.css">
                <script src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.5.0/highlight.min.js"></script>
                <style>
                    body {{ 
                        padding: 20px; 
                        font-family: Segoe UI; 
                        line-height: 1.6;
                        max-width: 800px;
                        margin: 0 auto;
                    }}
                    pre {{
                        background: #f8f9fa;
                        padding: 15px;
                        border-radius: 8px;
                        overflow-x: auto;
                    }}
                    .mathjax-block {{ 
                        padding: 10px;
                        margin: 20px 0;
                        background: #f3f4f6;
                        border-radius: 4px;
                    }}
                    code {{ font-family: 'Fira Code', monospace; }}
                </style>
                {mathjax_config}
            </head>
            <body>
                {html}
                <script>
                    document.addEventListener('DOMContentLoaded', (event) => {{
                        hljs.highlightAll();
                        if(typeof MathJax !== 'undefined') {{
                            MathJax.typesetPromise();
                        }}
                    }});
                </script>
            </body>
            </html>
        """)

    def show_tree_context_menu(self, pos):
        item = self.folder_tree.itemAt(pos)
        menu = QMenu()
        
        if item and item.parent():  # 笔记项
            edit_action = QAction("重命名", self)
            edit_action.triggered.connect(lambda: self.rename_note(item))
            delete_action = QAction("删除", self)
            delete_action.triggered.connect(lambda: self.delete_note(item))
            move_action = QAction("移动到...", self)
            move_action.triggered.connect(lambda: self.move_note(item))
            menu.addActions([edit_action, delete_action, move_action])
        elif item:  # 文件夹项
            rename_action = QAction("重命名文件夹", self)
            rename_action.triggered.connect(lambda: self.rename_folder(item))
            delete_action = QAction("删除文件夹", self)
            delete_action.triggered.connect(lambda: self.delete_folder(item))
            menu.addActions([rename_action, delete_action])
        else:  # 空白处
            new_folder_action = QAction("新建文件夹", self)
            new_folder_action.triggered.connect(self.create_folder)
            menu.addAction(new_folder_action)
            
        menu.exec_(self.folder_tree.viewport().mapToGlobal(pos))

    def show_tag_context_menu(self, pos):
        item = self.tag_list.itemAt(pos)
        menu = QMenu()
        if item and item.text() != "无标签":
            delete_action = QAction("删除标签", self)
            delete_action.triggered.connect(lambda: self.delete_tag(item))
            menu.addAction(delete_action)
        menu.exec_(self.tag_list.viewport().mapToGlobal(pos))

    def create_folder(self):
        folder_name, ok = QInputDialog.getText(
            self, "新建文件夹", "输入文件夹名称:",
            QLineEdit.Normal, ""
        )
        if ok and folder_name:
            # 检查是否已存在
            for folder in self.folders:
                if folder["name"] == folder_name:
                    QMessageBox.warning(self, "警告", "该文件夹已存在!")
                    return
            
            # 创建新文件夹
            new_folder = {
                "id": f"folder_{datetime.now().timestamp()}",
                "name": folder_name,
                "path": folder_name,
                "created": datetime.now().isoformat()
            }
            self.folders.append(new_folder)
            self.expanded_paths.add(folder_name)
            self.save_data()
            self.update_views()

    def rename_folder(self, item):
        old_name = item.text(0)
        new_name, ok = QInputDialog.getText(
            self, "重命名文件夹", "新名称:",
            QLineEdit.Normal, old_name
        )
        if ok and new_name and new_name != old_name:
            # 检查是否已存在
            for folder in self.folders:
                if folder["name"] == new_name:
                    QMessageBox.warning(self, "警告", "该名称已存在!")
                    return
            
            # 更新文件夹记录
            for folder in self.folders:
                if folder["name"] == old_name:
                    folder["name"] = new_name
                    folder["path"] = new_name
                    break
            
            # 更新笔记路径
            for note in self.notes:
                if note.get("path") == old_name:
                    note["path"] = new_name
            
            # 更新展开状态
            if old_name in self.expanded_paths:
                self.expanded_paths.remove(old_name)
                self.expanded_paths.add(new_name)
            
            self.save_data()
            self.update_views()

    def delete_folder(self, item):
        folder_name = item.text(0)
        confirm = QMessageBox.question(
            self, "删除确认",
            f"确定删除文件夹 '{folder_name}' 及其所有内容？",
            QMessageBox.Yes | QMessageBox.No
        )
        if confirm == QMessageBox.Yes:
            # 从文件夹列表中移除
            self.folders = [f for f in self.folders if f["name"] != folder_name]
            
            # 从展开状态中移除
            if folder_name in self.expanded_paths:
                self.expanded_paths.remove(folder_name)
            
            # 移动文件夹下的笔记到默认文件夹
            for note in self.notes:
                if note.get("path") == folder_name:
                    note["path"] = "默认文件夹"
            
            self.save_data()
            self.update_views()

    def move_note(self, item):
        note_id = item.data(0, Qt.UserRole)
        if not note_id:
            return
            
        # 获取可用文件夹列表
        folders = [f["name"] for f in self.folders]
        if not folders:
            QMessageBox.information(self, "提示", "没有可用的文件夹")
            return
            
        folder_name, ok = QInputDialog.getItem(
            self, "移动笔记", "选择目标文件夹:",
            folders, 0, False
        )
        
        if ok and folder_name:
            for note in self.notes:
                if note["id"] == note_id:
                    note["path"] = folder_name
                    break
            self.save_data()
            self.update_views()

    def delete_note(self, item):
        note_id = item.data(0, Qt.UserRole)
        confirm = QMessageBox.question(
            self, "删除确认", 
            "确定删除该笔记？此操作不可恢复！",
            QMessageBox.Yes | QMessageBox.No
        )
        if confirm == QMessageBox.Yes:
            # 删除笔记文件
            note_file = os.path.join(NOTES_DIR, f"{note_id}.txt")
            if os.path.exists(note_file):
                os.remove(note_file)
                
            # 从笔记列表中移除
            self.notes = [n for n in self.notes if n["id"] != note_id]
            
            self.save_data()
            self.update_views()

    def delete_tag(self, item):
        tag = item.text()
        if tag == "无标签":
            return
            
        # 从所有笔记中移除该标签
        for note in self.notes:
            if tag in note.get("tags", []):
                note["tags"].remove(tag)
        
        # 从标签列表移除
        self.tags.remove(tag)
        self.update_tag_list()
        self.save_data()

    def load_data(self):
        try:
            with open(NOTES_FILE, "r") as f:
                data = json.load(f)
                self.notes = data.get("notes", [])
                # 修复文件夹数据
                self.folders = []
                for folder in data.get("folders", []):
                    # 确保包含必要字段
                    self.folders.append({
                        "id": folder.get("id", f"folder_{datetime.now().timestamp()}"),
                        "name": folder.get("name", "未命名文件夹"),
                        "path": folder.get("path", folder.get("name", "未命名文件夹")),
                        "created": folder.get("created", datetime.now().isoformat())
                    })
                    self.expanded_paths = set(data.get("expanded_folders", []))
                
                # 初始化默认文件夹
                if not self.folders:
                    self.folders.append({
                        "id": "folder_default",
                        "name": "默认文件夹",
                        "path": "默认文件夹",
                        "created": datetime.now().isoformat()
                    })
                    self.expanded_paths.add("默认文件夹")
                
                # 加载标签
                all_tags = {tag for note in self.notes for tag in note.get("tags", [])}
                self.tags = ["所有标签", "无标签"] + sorted(all_tags)
                
                # 加载笔记内容
                for note in self.notes:
                    file_path = os.path.join(NOTES_DIR, f"{note['id']}.txt")
                    if os.path.exists(file_path):
                        with open(file_path, "r", encoding="utf-8") as f:
                            note["content"] = f.read()
        except (FileNotFoundError, json.JSONDecodeError):
            # 初始化默认数据
            self.notes = []
            self.folders = [{
                "id": "folder_default",
                "name": "默认文件夹",
                "path": "默认文件夹",
                "created": datetime.now().isoformat()
            }]
            self.expanded_paths = set(["默认文件夹"])
            self.save_data()

        self.update_views()

    def save_data(self):
        data = {
            "notes": [{
                "id": n["id"],
                "title": n["title"],
                "tags": n.get("tags", []),
                "path": n.get("path", "默认文件夹"),
                "created": n.get("created", ""),
                "modified": datetime.now().isoformat()
            } for n in self.notes],
            "folders": self.folders,
            "expanded_folders": list(self.expanded_paths)
        }
        
        with open(NOTES_FILE, "w") as f:
            json.dump(data, f, indent=2)
            
        # 保存笔记内容到单独文件
        for note in self.notes:
            file_path = os.path.join(NOTES_DIR, f"{note['id']}.txt")
            with open(file_path, "w", encoding="utf-8") as f:
                f.write(note.get("content", ""))
        
        self.content_updated.emit({"type": "notes", "data": data})

    def update_views(self):
        # 保存当前选中项
        current_item = self.folder_tree.currentItem()
        current_id = current_item.data(0, Qt.UserRole) if current_item else None
        
        self.folder_tree.clear()
        root = self.folder_tree.invisibleRootItem()
        
        # 构建文件夹树
        folder_items = {}
        for folder in sorted(self.folders, key=lambda x: x.get("name", "")):
            folder_item = QTreeWidgetItem(root, [folder["name"]])
            folder_item.setData(0, Qt.UserRole, folder["id"])
            folder_items[folder["name"]] = folder_item
            
            # 恢复展开状态
            if folder["name"] in self.expanded_paths:
                folder_item.setExpanded(True)
        
        # 添加笔记到对应文件夹
        for note in self.notes:
            folder_name = note.get("path", "默认文件夹")
            parent = folder_items.get(folder_name, root)
            
            note_item = QTreeWidgetItem(parent, [note["title"]])
            note_item.setData(0, Qt.UserRole, note["id"])
        
        # 恢复选中状态
        if current_id:
            iterator = QTreeWidgetItemIterator(self.folder_tree)
            while iterator.value():
                item = iterator.value()
                if item.data(0, Qt.UserRole) == current_id:
                    self.folder_tree.setCurrentItem(item)
                    break
                iterator += 1
        
        # 更新标签列表
        self.update_tag_list()

    def update_tag_list(self):
        self.tag_list.clear()
        all_tags = set()
        for note in self.notes:
            all_tags.update(note.get("tags", []))
        
        self.tags = ["所有标签", "无标签"] + sorted(all_tags)
        self.tag_list.addItems(self.tags)
        if self.tag_list.count() > 0:
            self.tag_list.item(0).setSelected(True)

    def create_note(self):
        # 确定默认文件夹
        default_folder = "默认文件夹"
        current_item = self.folder_tree.currentItem()
        if current_item and current_item.data(0, Qt.UserRole).startswith("folder_"):
            default_folder = current_item.text(0)
        
        note_id = f"note_{datetime.now().timestamp()}"
        self.current_note = {
            "id": note_id,
            "title": "新笔记",
            "content": "",
            "tags": [],
            "path": default_folder,
            "created": datetime.now().isoformat(),
            "modified": datetime.now().isoformat()
        }
        self.notes.append(self.current_note)
        self.save_data()
        self.update_views()
        self.load_note_data()

    def load_note(self, item):
        note_id = item.data(0, Qt.UserRole)
        if note_id and not note_id.startswith("folder_"):
            self.current_note = next((n for n in self.notes if n["id"] == note_id), None)
            self.load_note_data()

    def load_note_data(self):
        if self.current_note:
            self.title_input.setText(self.current_note.get("title", ""))
            self.tag_input.setText(", ".join(self.current_note.get("tags", [])))
            self.editor.setPlainText(self.current_note.get("content", ""))
            self.update_preview()

    def add_note(self, note_data):
        try:
            new_note = {
                "id": f"note_{datetime.now().timestamp()}",
                "title": note_data["title"],
                "content": note_data["content"],
                "tags": note_data.get("tags", []),
                "path": note_data.get("path", "默认文件夹"),
                "created": datetime.now().isoformat(),
                "modified": datetime.now().isoformat()
            }
            self.notes.append(new_note)
            self.save_data()
            self.update_views()

            # 定位到新笔记
            iterator = QTreeWidgetItemIterator(self.folder_tree)
            while iterator.value():
                item = iterator.value()
                if item.data(0, Qt.UserRole) == new_note["id"]:
                    self.folder_tree.setCurrentItem(item)
                    self.folder_tree.scrollToItem(item)
                    self.load_note(item)
                    break
                iterator += 1
        except KeyError as e:
            QMessageBox.critical(self, "错误", f"缺少必要字段: {str(e)}")

    def filter_by_tag(self, item):
        selected_tag = item.text()
        self.folder_tree.clear()
        root = QTreeWidgetItem(self.folder_tree, [f"标签: {selected_tag}"])
        
        if selected_tag == "所有标签":
            filtered = [n for n in self.notes if n.get("tags")]
        elif selected_tag == "无标签":
            filtered = [n for n in self.notes if not n.get("tags")]
        else:
            filtered = [n for n in self.notes if selected_tag in n.get("tags", [])]
            
        for note in filtered:
            item = QTreeWidgetItem(root, [note["title"]])
            item.setData(0, Qt.UserRole, note["id"])
            
        root.setExpanded(True)

    def save_current(self):
        if self.current_note:
            self.current_note.update({
                "title": self.title_input.text(),
                "tags": [t.strip() for t in self.tag_input.text().split(",") if t.strip()],
                "content": self.editor.toPlainText(),
                "modified": datetime.now().isoformat()
            })
            
            self.save_data()
            self.update_views()
            
            # 定位到当前笔记
            iterator = QTreeWidgetItemIterator(self.folder_tree)
            while iterator.value():
                item = iterator.value()
                if item.data(0, Qt.UserRole) == self.current_note["id"]:
                    self.folder_tree.setCurrentItem(item)
                    break
                iterator += 1

    def get_item_path(self, item):
        """获取树项的完整路径"""
        path = []
        while item and item.parent():
            path.append(item.text(0))
            item = item.parent()
        return "/".join(reversed(path))